{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ac35e0b4",
   "metadata": {},
   "source": [
    "# NumPy-like Operations\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/brainpy/blob/master/docs_version2/tutorial_math/Numpy_like_Operations.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/brainpy/blob/master/docs_version2/tutorial_math/Numpy_like_Operations.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e856b15",
   "metadata": {},
   "source": [
    "@[Xiaoyu Chen](mailto:c-xy17@tsinghua.org.cn)\n",
    "@[Chaoming Wang](https://github.com/chaoming0625)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ccd2442",
   "metadata": {},
   "source": [
    "## `brainpy.math.array`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "195e295f",
   "metadata": {},
   "source": [
    "Neural modeling often involves the computation of thousands of elements and requires large-scale computation. To support this, BrainPy provides a basic data structure `brainpy.math.array` with abundant operations. In fact, the concept of `brainpy.math.array` is the same as `ndarray` in NumPy, and the operations for `brainpy.math.array` closely resemble NumPy operations. \n",
    "\n",
    "\n",
    "Users who are not familiar with the NumPy package may refer to [https://numpy.org/](https://numpy.org/) for more information.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d2fa1cd",
   "metadata": {},
   "source": [
    "## NumPy-like operations for arrays"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab12bd3c",
   "metadata": {},
   "source": [
    "For `brainpy.math.array`, BrainPy provides array-related operations that have the same form as NumPy while capable of performing JIT transformations with the help of JAX. Almost all NumPy operations can be used in BrainPy as long as users transform `numpy.xxx` to `brainpy.math.xxx`.\n",
    "\n",
    "Specifically, the NumPy-like operations for `brainpy.math.array` include:\n",
    "\n",
    "1. **Creating an array**: users can use `brainpy.math.array()`， `brainpy.math.ones()`， `brainpy.math.zeros()`， `brainpy.math.arange()`， and so on to create an array.\n",
    "2. **Element-wise operations**: `+`, `-`, `*`, `/`, ...\n",
    "3. **Aggregation functions**: `.max()`, `.min()`, `.sum()`, `.mean()`, `.prod()`, ...\n",
    "4. **Broadcasting**: operating two arrays that have different dimensions (usually a smaller array is broadcast to a larger array with compatible shapes to complete operation).\n",
    "5. **Indexing, slicing, and iterating**: Though based on JAX, BrainPy provides NumPy-like indexing, slicing, and iterating operations that are more convenient than those in JAX. Arrays can be treated much the same as other Python sequence data structures such as lists.\n",
    "6. **Random functions**: BrainPy also offers NumPy-like random functions that are different from JAX. To enable random functions in JIT compilation, JAX requires users to pass an explicit key to define a random state, whereas random functions in BrainPy can be used in the same way as in NumPy where the seed does not need to be explicitly defined.\n",
    "7. **Other mathematical and logical functions**: `brainpy.math.sin()`, `brainpy.math.sort()`, `brainpy.math.argmax()`, `brainpy.math.where()`, ..."
   ]
  },
  {
   "cell_type": "code",
   "id": "ee870ad4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:23.315768Z",
     "start_time": "2025-10-06T04:02:19.032952Z"
    }
   },
   "source": [
    "import brainpy.math as bm\n",
    "\n",
    "bm.set_platform('cpu')"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "id": "01f00053",
   "metadata": {},
   "source": [
    "Here are some examples of array-related operations in BrainPy:"
   ]
  },
  {
   "cell_type": "code",
   "id": "d180ee8d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:23.341996Z",
     "start_time": "2025-10-06T04:02:23.320290Z"
    }
   },
   "source": [
    "# to create an array\n",
    "a1 = bm.array([[0, 1, 2], [3, 4, 5]])\n",
    "a1"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([[0, 1, 2],\n",
       "                   [3, 4, 5]]),\n",
       "      dtype=int32)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "id": "bfbcfa02",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:23.387087Z",
     "start_time": "2025-10-06T04:02:23.350558Z"
    }
   },
   "source": [
    "# element-wise operation\n",
    "a1 += 1\n",
    "a1"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([[1, 2, 3],\n",
       "                   [4, 5, 6]]),\n",
       "      dtype=int32)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 3
  },
  {
   "cell_type": "code",
   "id": "bbe6fff9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:23.425065Z",
     "start_time": "2025-10-06T04:02:23.392104Z"
    }
   },
   "source": [
    "a2 = bm.ones(3)\n",
    "a2"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([1., 1., 1.]), dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 4
  },
  {
   "cell_type": "code",
   "id": "bdfae8f6",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:23.586249Z",
     "start_time": "2025-10-06T04:02:23.438487Z"
    }
   },
   "source": [
    "# slicing and indexing\n",
    "a2[0:2] = 10.\n",
    "a2[2] = 0.\n",
    "a2"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([10., 10.,  0.]), dtype=float32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 5
  },
  {
   "cell_type": "code",
   "id": "9a470aea",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:23.641207Z",
     "start_time": "2025-10-06T04:02:23.607601Z"
    }
   },
   "source": [
    "# broadcasting\n",
    "a3 = a1 + a2\n",
    "a3"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[11., 12.,  3.],\n",
       "       [14., 15.,  6.]], dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "cell_type": "code",
   "id": "13d4609f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:23.671101Z",
     "start_time": "2025-10-06T04:02:23.645547Z"
    }
   },
   "source": [
    "# sorting\n",
    "bm.sort(a3, axis=-1)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([[ 3., 11., 12.],\n",
       "                   [ 6., 14., 15.]]),\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 7
  },
  {
   "cell_type": "code",
   "id": "e1af1f2a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:23.718009Z",
     "start_time": "2025-10-06T04:02:23.687630Z"
    }
   },
   "source": [
    "# reshaping\n",
    "a1.reshape(6, -1)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[1],\n",
       "       [2],\n",
       "       [3],\n",
       "       [4],\n",
       "       [5],\n",
       "       [6]], dtype=int32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "id": "2535f8ed",
   "metadata": {},
   "source": [
    "Here is an example for random number generation in JAX and in BrainPy:"
   ]
  },
  {
   "cell_type": "code",
   "id": "1b781ee4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:24.015548Z",
     "start_time": "2025-10-06T04:02:23.725975Z"
    }
   },
   "source": [
    "# random number generation in JAX\n",
    "from jax import random\n",
    "\n",
    "key = random.PRNGKey(0)  # define a key explicitly\n",
    "r1 = random.uniform(key)  # a random number generated from a uniform distribution\n",
    "r2 = random.uniform(key)  # reusing the same key leads to the same result\n",
    "print(r1 == r2)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "execution_count": 9
  },
  {
   "cell_type": "code",
   "id": "700a6610",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:24.136424Z",
     "start_time": "2025-10-06T04:02:24.017068Z"
    }
   },
   "source": [
    "# random number generation in BrainPy\n",
    "r3 = bm.random.uniform()  # the seed does not need to be passed explicitly\n",
    "r4 = bm.random.uniform(key=1)  # users can also assign an explicit key to generate a constant result\n",
    "r5 = bm.random.uniform(key=1)\n",
    "print(r3)\n",
    "print(r4 == r5)"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.050184727\n",
      "True\n"
     ]
    }
   ],
   "execution_count": 10
  },
  {
   "cell_type": "markdown",
   "id": "d7aec642",
   "metadata": {},
   "source": [
    "Users can also generate a random array using random functions in BrainPy:"
   ]
  },
  {
   "cell_type": "code",
   "id": "2ee263d5",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:24.270312Z",
     "start_time": "2025-10-06T04:02:24.164416Z"
    }
   },
   "source": [
    "bm.random.random((5, 3))  # to generate random floats number in the interval [0.0, 1.0)"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[0.59585726, 0.2890029 , 0.111938  ],\n",
       "       [0.5041219 , 0.15522647, 0.7115699 ],\n",
       "       [0.9315944 , 0.22322583, 0.38390863],\n",
       "       [0.6865432 , 0.86952937, 0.70655596],\n",
       "       [0.0250324 , 0.9579313 , 0.54371226]], dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 11
  },
  {
   "cell_type": "markdown",
   "id": "d6c4a6f3",
   "metadata": {},
   "source": [
    "For all array-related operations, please refer to the [`brianpy.math` module](../apis/auto/math.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27865417",
   "metadata": {},
   "source": [
    "## When to use `brainpy.math.array`?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "897b32b7",
   "metadata": {},
   "source": [
    "Simply speaking, users can use `brainpy.math.array` whenever array computation is involved in JIT compilation. While JIT compilation speeds up code execution, it also brings some limitations. Once an array is given to the JIT compiler, the values inside the array cannot be changed. This means `brainpy.math.array` can only be used to store values that do not change over time. Static neural properties such as the membrane capacitance, the voltage threshold, and the time constant of a neuron population, and static synaptic connections between neurons are good candidates to be defined as arrays (or a unified single value).\n",
    "\n",
    "To extend arrays to store changing values, BrainPy offers a new data structure named ``brainpy.math.Variable``. More details about Variables can be found in [`Variable` and `BrainPbject`](variables.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "id": "2b395ef1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T04:02:24.288242Z",
     "start_time": "2025-10-06T04:02:24.285326Z"
    }
   },
   "source": [],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
